package com.offcn.bigdata.spark.p4

import java.net.URL

import org.apache.spark.rdd.RDD
import org.apache.spark.{Partitioner, SparkConf, SparkContext}

import scala.collection.mutable

/**
  * 分组分区
  * 数据：
  *     20181123101523 http://java.ujiuye.com/java/javaee.shtml?kw=
        20181123101523 http://h5.ujiuye.com/java/javaee.shtml
    需求，按照不同的学科，将用户的访问信息进行输出保存，一个学科一个文件
    一个输出文件，代表的就是一个partition，也就需要按照访问特点进行分区
  */
object _02GroupedPartitionerOps {
    def main(args: Array[String]): Unit = {
        val conf = new SparkConf()
            .setMaster("local[*]")
            .setAppName(s"${_02GroupedPartitionerOps.getClass.getSimpleName}")
        val sc = new SparkContext(conf)

        val lines = sc.textFile("file:/E:/work/2020-0828期大数据/workspace/spark-parent-0828/data/access.txt")

        val pairs:RDD[(String, String)] = lines.map(line => {
            val fields = line.split("\\s+")
            if(fields == null || fields.length != 2) {
                null
            } else {
                (fields(1).trim, line)
            }
        }).filter(_ != null)

        val keys = pairs.keys.collect().toSet
        val hosts = keys.map(line => {
            val url = new URL(line)
            url.getHost
        }).toArray
        println(s"抽取出来的学科：${hosts.mkString("[", ", ", "]")}")

        pairs.partitionBy(new GroupedPartitioner(hosts))
            .values
            .saveAsTextFile("file:/E:/data/out/grouped")

        sc.stop()
    }
}
/*
    通过host去映射partition编号。
        在自定义分区中传入提取出来的学科信息，完成组名称到分区编号的映射。
        在getPartition中，通过key提取学科，然后再从映射表中获取分区编号
 */
class GroupedPartitioner(hosts: Array[String]) extends Partitioner {
    override def numPartitions: Int = hosts.length
    val host2Partition = {
        val map = mutable.HashMap[String, Int]()
        for(i <- 0 until hosts.length) {
            map.put(hosts(i), i)
        }
        map
    }

    override def getPartition(key: Any): Int = {
        val host = new URL(key.toString).getHost
        host2Partition.getOrElse(host, 0)
    }
}
